#include "nBodyHolder.h"

// Default constructor. Creates an empty n-body holder.
nBodyHolder::nBodyHolder()
{
   nBodies = 0;
   head = NULL;
}

// Creates an n-body holder by reading in a file.
nBodyHolder::nBodyHolder(string fn, ParamHolder* par)
{
   readFromFile(fn);
   params = par;
}

// Reads in a file containing the bodies.
// The format is "mass, pos[X], pos[y], pos[z]
// If the holder is not empty, anything currently in it gets erased.
void nBodyHolder::readFromFile(string fn)
{
   nBodies = 0;
   head = NULL;

   ifstream in;
   in.open(fn.c_str());

   string line;
   double nm, nx, ny, nz, nvx, nvy, nvz;

   while (in.good()) {
      getline(in, line);
      line = prepLine(line);

      if (line.length() > 0) {
         nm = getNum(line);

	 line = rmFirstNum(line);

         nx = getNum(line);
	 line = rmFirstNum(line);

         ny = getNum(line);
	 line = rmFirstNum(line);

         nz = getNum(line);
	 line = rmFirstNum(line);

         nvx = getNum(line);
	 line = rmFirstNum(line);

         nvy = getNum(line);
	 line = rmFirstNum(line);

         nvz = getNum(line);
	 line = rmFirstNum(line);

         vec nPos = vec(3);
         vec nVel = vec(3);

	 nPos[X] = nx;
	 nPos[Y] = ny;
	 nPos[Z] = nz;

	 nVel[X] = nvx;
	 nVel[Y] = nvy;
	 nVel[Z] = nvz;

	 Body* b = new Body(nm, nPos, nVel);

         addBody(b);
	 }
    } // End loop over lines.

    centerSystem();
}

// Writes the current state of the simulation to a text file.
// The text file is comma delimited and human readable, but can
// also be used to initialize the simulation.
//
// The format is "mass, pos[X], pos[y], pos[z]
void nBodyHolder::writeToFile(string fn)
{
   ofstream of;
   of.open(fn.c_str());
   
   Body* b = NULL;

   // The width of the output should be a parameter.
   for(int i = 0; i < nBodies; i++) {
      b = getBody(i);

      of << setw(15) << b -> mass << ",";

      of << setw(15) << b -> pos[X] << ",";
      of << setw(15) << b -> pos[Y] << ",";
      of << setw(15) << b -> pos[Z] << ",";

      of << setw(15) << b -> vel[X] << ",";
      of << setw(15) << b -> vel[Y] << ",";
      of << setw(15) << b -> vel[Z] << endl;
      }

   of.close();
}

// Currently, everything is stored in a linked list.
// No sort of order is gaurenteed.
// Everything is added to the front of the list, meaning that adding
// a body changes the indexing.
void nBodyHolder::addBody(Body* nb)
{
   // Add to the beginning of the list
   BodyNode* newHead = new BodyNode();
   newHead -> b = nb;
   newHead -> next = head;

   head = newHead;

   nBodies ++;
}

// Treating the holder as if it were an array, returns the body at
// index i. No sort of ordering is guarenteed, but it is guarenteed
// that each index will return a different body and that all bodies
// will have an index between 0 and nBodies-1.
Body* nBodyHolder::getBody(int i)
{
   BodyNode* curr = head;

   int j = 0;

   if (i < nBodies) {
      while (j != i) {
         curr = curr -> next;
	 j++;
	 }

      return curr -> b;
      }

   else
      return NULL; // Index exceeds size. Should throw an error.
}

// Accessor function for the number of bodies in the simulation.
int nBodyHolder::size()
{
   return nBodies;
}

// Compute the acceleration on body j.
vec nBodyHolder::accel(int j)
{
   vec a = vec(3);
   vec netForce = vec(3);

   for(int i = 0; i < 3; i++) 
      netForce[i] = 0.0;

   Body* m1 = getBody(j);
   Body* m2 = NULL;
   double dist, force;
   double G = params -> G_cgs / (pow(params -> LengthUnits, 3.0)
              / (params -> MassUnits * params -> TimeUnits
	      * params -> TimeUnits));

   for (int i = 0; i < nBodies; i++) {
      if (i != j) {
         m2 = getBody(i);

         dist = sqrt(pow(m1->pos[X] - m2->pos[X], 2.0)
                   + pow(m1->pos[Y] - m2->pos[Y], 2.0)
	           + pow(m1->pos[Z] - m2->pos[Z], 2.0));

         if (dist < 0.01)
	    dist = 0.01;

         double eps = params -> SofteningParameter;

         force = G * m1->mass * m2->mass / (dist*dist + eps * eps);

         netForce[X] -= force * (m1->pos[X] - m2->pos[X]) / dist;
         netForce[Y] -= force * (m1->pos[Y] - m2->pos[Y]) / dist;
         netForce[Z] -= force * (m1->pos[Z] - m2->pos[Z]) / dist;
         }
      }

   a[X] = netForce[X]/m1->mass;
   a[Y] = netForce[Y]/m1->mass;
   a[Z] = netForce[Z]/m1->mass;

   return a;
}

// Removes comments and spaces
string nBodyHolder::prepLine(string line) {
   // Remove comments.
   int pos = line.find("#");

   if (pos != string::npos)
   line = line.substr(0, pos);

   pos = line.find("//");

   if (pos != string::npos)
   line = line.substr(0, pos);

   // Remove spaces
   for (int i = 0; i < line.length(); i++) {
      if (line[i] == ' ')
         line = line.substr(0, i) + line.substr(i+1, line.length());
	    }

   return line;
}

// Returns the first number in a comma delimited string.
double nBodyHolder::getNum(string s)
{
   int pos = s.find_first_of(',');
   string numstr = s.substr(0, pos);

   istringstream ss(numstr);

   float num;
   ss >> num;

   return (double)num;
}

// Removes the first number in a comma delimited string. Also removes
// the comma.
string nBodyHolder::rmFirstNum(string s)
{
   int pos = s.find(',');
   return s.substr(pos + 1, s.size() - pos);
}

// Adjust the positions and velocitie so that the center of mass is
// fixed at the center of the system.
void nBodyHolder::centerSystem()
{
   vec cm = vec(3); // Compute the position of the center of mass.
   vec cv = vec(3); // Compute the velocity of the center of mass.

   double totalMass = 0.0;

   int n = nBodies;

   for (int i = 0; i < n; i++) {
      Body* b = getBody(i);

      cm[X] += b -> mass * b -> pos[X];
      cm[Y] += b -> mass * b -> pos[Y];
      cm[Z] += b -> mass * b -> pos[Z];

      cv[X] += b -> mass * b -> vel[X];
      cv[Y] += b -> mass * b -> vel[Y];
      cv[Z] += b -> mass * b -> vel[Z];

      totalMass += b -> mass;
      }

   cm[X] /= totalMass;
   cm[Y] /= totalMass;
   cm[Z] /= totalMass;

   cv[X] /= totalMass;
   cv[Y] /= totalMass;
   cv[Z] /= totalMass;

   for (int i = 0; i < n; i++) {
      Body* b = getBody(i);

      b -> pos[X] -= cm[X];
      b -> pos[Y] -= cm[Y];
      b -> pos[Z] -= cm[Z];

      b -> vel[X] -= cv[X];
      b -> vel[Y] -= cv[Y];
      b -> vel[Z] -= cv[Z];
      }
}
